import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
from collections import deque
import time

# 1 决策网络构建
# 1.1 干扰带宽决策网络
class BandDecisionNet(nn.Module):
    def __init__(self, input_dim):
        super(BandDecisionNet, self).__init__()
        self.fc1 = nn.Linear(input_dim, 16)
        self.fc2 = nn.Linear(16, 16)
        self.fc3 = nn.Linear(16, 1)  # 输出带宽选择（连续值）
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 1.2 干扰频段决策网络
class FreqDecisionNet(nn.Module):
    def __init__(self, input_dim):
        super(FreqDecisionNet, self).__init__()
        hidden_size = input_dim * 2  # 隐藏层神经元数量根据输入维度动态计算
        self.fc1 = nn.Linear(input_dim, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, 1)  # 输出频段选择（连续值）
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 2 经验回放缓冲区
class BETMR:
    def __init__(self, capacity, expert_capacity):
        self.capacity = capacity
        self.expert_capacity = expert_capacity
        self.normal_buffer = deque(maxlen=capacity)
        self.expert_buffer = deque(maxlen=expert_capacity)
        self.thresholds = [0, 10, 20, 30]  # 动态阈值集
        self.current_threshold_index = 0

    def push(self, state, action, reward, next_state, done):
        self.normal_buffer.append((state, action, reward, next_state, done))

    def store_expert(self, state, action, reward, next_state, done):
        self.expert_buffer.append((state, action, reward, next_state, done))

    def update_threshold(self, total_reward):
        if self.current_threshold_index < len(self.thresholds) - 1 and total_reward > self.thresholds[
            self.current_threshold_index + 1]:
            self.current_threshold_index += 1
        elif self.current_threshold_index > 0 and total_reward < self.thresholds[self.current_threshold_index - 1]:
            self.current_threshold_index -= 1
        # 如果达到最高阈值，重置专家经验池
        if self.current_threshold_index == len(self.thresholds) - 1:
            self.expert_buffer.clear()
            self.current_threshold_index = 0

    # def sample(self, batch_size):
    #     if not self.normal_buffer and not self.expert_buffer:
    #         return None, None
    #     # 按照公式 (7) 抽样
    #     if self.expert_buffer:
    #         expert_ratio = 0.3  # 专家样本占30%
    #         expert_batch = min(int(batch_size * expert_ratio), len(self.expert_buffer))
    #         normal_batch = batch_size - expert_batch
    #         expert_samples = random.sample(self.expert_buffer, expert_batch)
    #         normal_samples = random.sample(self.normal_buffer, normal_batch)
    #         return expert_samples + normal_samples
    #     else:
    #         return random.sample(self.normal_buffer, batch_size)
    def sample(self, batch_size):
        if not self.normal_buffer and not self.expert_buffer:
            return None, None
        # 按照新阈值抽样
        if self.expert_buffer:
            expert_ratio = 0.3  # 专家样本占30%
            expert_batch = min(int(batch_size * expert_ratio), len(self.expert_buffer))
            normal_batch = batch_size - expert_batch
            expert_samples = random.sample(self.expert_buffer, expert_batch)
            normal_samples = random.sample(self.normal_buffer, normal_batch)
            return expert_samples + normal_samples
        else:
            return random.sample(self.normal_buffer, batch_size)

# 3 BHJM 分层强化学习模型
class BHJM:
    def __init__(self, state_dim, lr=1e-4, gamma=0.99, buffer_capacity=10000, expert_capacity=2000, batch_size=128):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")

        self.freq_net = FreqDecisionNet(state_dim).to(self.device)
        self.band_net = BandDecisionNet(state_dim).to(self.device)

        self.freq_target_net = FreqDecisionNet(state_dim).to(self.device)
        self.band_target_net = BandDecisionNet(state_dim).to(self.device)

        self.freq_target_net.load_state_dict(self.freq_net.state_dict())
        self.band_target_net.load_state_dict(self.band_net.state_dict())
        self.freq_target_net.eval()
        self.band_target_net.eval()

        self.freq_optimizer = optim.Adam(self.freq_net.parameters(), lr=lr)
        self.band_optimizer = optim.Adam(self.band_net.parameters(), lr=lr)

        self.gamma = gamma
        self.betmr = BETMR(buffer_capacity, expert_capacity)
        self.batch_size = batch_size
        self.criterion = nn.MSELoss()
        self.tau = 0.001
        self.thresholds = [0, 10, 20, 30]
        self.current_episode = 0
        self.reward_history = []  # 奖励归一化所需

    def update_threshold(self, total_reward):
        self.current_episode += 1

        # 每 100 轮动态调整阈值
        if self.current_episode % 100 == 0 and len(self.reward_history) > 100:
            mean = np.mean(self.reward_history[-100:])  # 取最近100个奖励的均值
            std = np.std(self.reward_history[-100:])
            self.thresholds = [
                0,
                mean,
                mean + std,
                mean + 2 * std
            ]

    def choose_action(self, state, epsilon=0.1):
        state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0).to(self.device)
        with torch.no_grad():
            freq = self.freq_net(state_tensor)
            band = self.band_net(state_tensor)
        # ε-greedy 策略
        if np.random.rand() < epsilon:
            freq = torch.tensor([np.random.uniform(-1, 1)], device=self.device)
            band = torch.tensor([np.random.uniform(0.5, 2.5)], device=self.device)
        return freq.item(), band.item()

    def store_transition(self, state, action, reward, next_state, done):
        self.betmr.push(state, action, reward, next_state, done)

    def update(self):
        if len(self.betmr.normal_buffer) < self.batch_size:
            return

        samples = self.betmr.sample(self.batch_size)
        if not samples:
            return

        states = []
        actions = []
        rewards = []
        next_states = []
        dones = []
        for sample in samples:
            s, a, r, ns, d = sample
            states.append(s)
            actions.append(a)
            rewards.append(r)
            next_states.append(ns)
            dones.append(d)

        states = torch.tensor(np.array(states), dtype=torch.float32).to(self.device)
        actions = torch.tensor(np.array(actions), dtype=torch.float32).to(self.device)
        rewards = torch.tensor(np.array(rewards), dtype=torch.float32).unsqueeze(1).to(self.device)
        next_states = torch.tensor(np.array(next_states), dtype=torch.float32).to(self.device)
        dones = torch.tensor(np.array(dones), dtype=torch.float32).unsqueeze(1).to(self.device)

        # 当前Q值
        current_freq_q = self.freq_net(states)
        current_band_q = self.band_net(states)

        # 目标Q值（使用目标网络）
        with torch.no_grad():
            next_freq_q = self.freq_target_net(next_states)
            next_band_q = self.band_target_net(next_states)
            target_freq_q = rewards + (1 - dones) * self.gamma * next_freq_q
            target_band_q = rewards + (1 - dones) * self.gamma * next_band_q

        # 损失函数
        loss_freq = self.criterion(current_freq_q, target_freq_q)
        loss_band = self.criterion(current_band_q, target_band_q)

        # 更新主网络
        self.freq_optimizer.zero_grad()
        loss_freq.backward()
        self.freq_optimizer.step()

        self.band_optimizer.zero_grad()
        loss_band.backward()
        self.band_optimizer.step()

        # 软更新目标网络
        for target_param, source_param in zip(self.freq_target_net.parameters(), self.freq_net.parameters()):
            target_param.data.copy_(self.tau * source_param.data + (1.0 - self.tau) * target_param.data)
        for target_param, source_param in zip(self.band_target_net.parameters(), self.band_net.parameters()):
            target_param.data.copy_(self.tau * source_param.data + (1.0 - self.tau) * target_param.data)

    def normalize_reward(self, reward):
        """ 对奖励进行归一化 """
        self.reward_history.append(reward)
        if len(self.reward_history) > 100:  # 限制历史长度，防止内存爆炸
            self.reward_history.pop(0)
        self.mean_reward = np.mean(self.reward_history)
        self.std_reward = np.std(self.reward_history) + 1e-8  # 防止除零
        return (reward - self.mean_reward) / self.std_reward

# 4、模拟环境
class JammingEnv:
    def __init__(self):
        self.state_dim = 10
        self.max_freq = 400
        self.min_freq = 200
        self.max_bandwidth = 3

    def reset(self):
        self.state = np.random.rand(self.state_dim)
        return self.state

    def step(self, action):
        freq, band = action
        # 新奖励函数：非负，鼓励高频率和带宽
        reward = np.sin(freq * 0.1) + band * 0.5
        # 确保奖励非负
        reward = max(reward, 0)
        next_state = self.state + np.random.normal(0, 0.1, self.state_dim)
        next_state = np.clip(next_state, 0, 1)
        self.state = next_state
        done = False
        return next_state, reward, done

# 5 训练流程
def train():
    env = JammingEnv()
    state_dim = env.state_dim
    bhjm = BHJM(state_dim)

    episodes = 1000
    max_steps = 100
    epsilon = 0.2  # 初期探索率

    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        for step in range(max_steps):
            action = bhjm.choose_action(state, epsilon)
            next_state, reward, done = env.step(action)

            # 归一化奖励
            normalized_reward = bhjm.normalize_reward(reward)

            # 存储经验
            bhjm.store_transition(state, action, normalized_reward, next_state, done)
            state = next_state
            total_reward += reward
            if done:
                break
        # 更新阈值
        bhjm.update_threshold(total_reward)
        # 更新网络
        bhjm.update()
        # 逐步降低探索率
        epsilon = max(0.01, epsilon * 0.995)
        print(f"Episode {episode + 1}/{episodes}, Total Reward: {total_reward:.2f}")

# 6、运行代码
if __name__ == "__main__":
    train()